module conv #(
    parameter int unsigned P_ICH  = 4,
    parameter int unsigned P_OCH  = 4,
    parameter int unsigned N_ICH  = 16,
    parameter int unsigned N_OCH  = 16,
    parameter int unsigned K      = 3,
    parameter int unsigned A_BIT  = 8,
    parameter int unsigned W_BIT  = 8,
    parameter int unsigned B_BIT  = 32,
    parameter int unsigned N_HW   = 64,
    parameter string       W_FILE = ""
) (
    input  logic                   clk,
    input  logic                   rst_n,
    input  logic [P_ICH*A_BIT-1:0] in_data,
    input  logic                   in_valid,
    output logic                   in_ready,
    output logic [P_OCH*B_BIT-1:0] out_data,
    output logic                   out_valid,
    input  logic                   out_ready
);

    localparam int unsigned FOLD_I = N_ICH / P_ICH;
    localparam int unsigned FOLD_O = N_OCH / P_OCH;
    localparam int unsigned KK = K * K;
    localparam int unsigned WEIGHT_DEPTH = FOLD_O * FOLD_I * KK;

    logic [$clog2(WEIGHT_DEPTH)-1:0] weight_addr_d0;
    logic [   P_OCH*P_ICH*W_BIT-1:0] weight_data_d1;
    rom #(
        .DWIDTH(P_OCH * P_ICH * W_BIT),
        .AWIDTH($clog2(WEIGHT_DEPTH)),
        .MEM_SIZE(WEIGHT_DEPTH),
        .INIT_FILE(W_FILE)
    ) u_weight_rom (
        .clk  (clk),
        .ce0  (out_ready),
        .addr0(weight_addr_d0),
        .q0   (weight_data_d1)
    );

    logic        [     P_ICH*A_BIT-1:0] line_buffer           [FOLD_I] [KK];
    logic signed [           B_BIT-1:0] acc                   [ P_OCH];
    logic        [  $clog2(N_HW+1)-1:0] cntr_hw;
    logic        [$clog2(FOLD_O+1)-1:0] cntr_fo;
    logic        [$clog2(FOLD_I+1)-1:0] cntr_fi;
    logic        [    $clog2(KK+1)-1:0] cntr_kk;
    logic                               pipe_en_in;
    logic                               pipe_en_out;
    logic                               is_fst_fo_d0;
    logic                               mac_array_data_vld_d1;
    logic        [     P_ICH*A_BIT-1:0] in_buf_d1;
    logic                               is_fst_kk_fi_d1;
    logic                               is_lst_kk_fi_d1;
    logic                               is_lst_kk_fi_dly      [ P_ICH];

    assign is_fst_fo_d0   = (cntr_fo == 0);
    assign pipe_en_in     = is_fst_fo_d0 ? in_valid : 1'b1;
    assign pipe_en_out    = out_ready;
    assign pipe_en        = pipe_en_in && pipe_en_out;
    assign in_ready       = is_fst_fo_d0 && out_ready;
    assign weight_addr_d0 = (cntr_fo * KK * FOLD_I) + cntr_fi * KK + cntr_kk;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            cntr_hw <= 0;
            cntr_fo <= 0;
            cntr_fi <= 0;
            cntr_kk <= 0;
            for (int i = 0; i < FOLD_I; i++) begin
                for (int k = 0; k < KK; k++) begin
                    line_buffer[i][k] <= '0;
                end
            end
        end else if (pipe_en) begin
            if (is_fst_fo_d0) begin
                line_buffer[cntr_fi][cntr_kk] <= in_data;
            end
            if (cntr_kk == KK - 1) begin
                cntr_kk <= 0;
                if (cntr_fi == FOLD_I - 1) begin
                    cntr_fi <= 0;
                    if (cntr_fo == FOLD_O - 1) begin
                        cntr_fo <= 0;
                        if (cntr_hw == N_HW - 1) begin
                            cntr_hw <= 0;
                        end else begin
                            cntr_hw <= cntr_hw + 1;
                        end
                    end else begin
                        cntr_fo <= cntr_fo + 1;
                    end
                end else begin
                    cntr_fi <= cntr_fi + 1;
                end
            end else begin
                cntr_kk <= cntr_kk + 1;
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            in_buf_d1             <= '0;
            is_fst_kk_fi_d1       <= 1'b0;
            is_lst_kk_fi_d1       <= 1'b0;
            mac_array_data_vld_d1 <= 1'b0;
        end else if (pipe_en_out) begin
            in_buf_d1             <= is_fst_fo_d0 ? in_data : line_buffer[cntr_fi][cntr_kk];
            is_fst_kk_fi_d1       <= (cntr_kk == 0) && (cntr_fi == 0) && pipe_en_in;
            is_lst_kk_fi_d1       <= (cntr_kk == KK - 1) && (cntr_fi == FOLD_I - 1) && pipe_en_in;
            mac_array_data_vld_d1 <= (is_fst_fo_d0 ? in_valid : 1'b1);
        end
    end

    logic        [A_BIT-1:0] x_vec[P_ICH];
    logic signed [W_BIT-1:0] w_vec[P_OCH] [P_ICH];
    always_comb begin
        for (int i = 0; i < P_ICH; i++) begin
            x_vec[i] = in_buf_d1[i*A_BIT+:A_BIT];
        end
    end

    always_comb begin
        for (int o = 0; o < P_OCH; o++) begin
            for (int i = 0; i < P_ICH; i++) begin
                w_vec[o][i] = weight_data_d1[(P_ICH*o+i)*W_BIT+:W_BIT];
            end
        end
    end

    generate
        for (genvar o = 0; o < P_OCH; o++) begin : gen_mac_array
            mac_array #(
                .P_ICH(P_ICH),
                .A_BIT(A_BIT),
                .W_BIT(W_BIT),
                .B_BIT(B_BIT)
            ) u_mac_array (
                .clk    (clk),
                .rst_n  (rst_n),
                .en     (pipe_en_out),
                .dat_vld(mac_array_data_vld_d1),
                .clr    (is_fst_kk_fi_d1),
                .x_vec  (x_vec),
                .w_vec  (w_vec[o]),
                .acc    (acc[o])
            );
        end
    endgenerate

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            for (int i = 0; i < P_ICH; i++) begin
                is_lst_kk_fi_dly[i] <= 1'b0;
            end
        end else if (pipe_en_out) begin
            is_lst_kk_fi_dly[0] <= is_lst_kk_fi_d1;
            for (int i = 1; i < P_ICH; i++) begin
                is_lst_kk_fi_dly[i] <= is_lst_kk_fi_dly[i-1];
            end
        end
    end

    assign out_valid = is_lst_kk_fi_dly[P_ICH-1];

    always_comb begin
        for (int o = 0; o < P_OCH; o++) begin
            out_data[o*B_BIT+:B_BIT] = acc[o];
        end
    end
endmodule
